//
// Created by Orlando Chen on 2023/05/03.
//

#include "../linux_nl/common.h"

#if __linux__
#include <arpa/inet.h>
#include <sys/socket.h>
#include <linux/netlink.h>
#include <cstring>
#include <csignal>
#include <unistd.h>
#include <cstdio>
#include <cstdlib>
#include <fcntl.h>
#else
// As Apple does not support netlink, when the user compiles this program on macOS,
// this program will prompt the user to compile the program on a Linux machine.
#error "This program is not supported on macOS. Please compile this program on a Linux machine."
#endif


int create_tcp_client_socket(const std::string& server_ip, int port) {

    /* Local variables*/
    int client_fd = socket(AF_INET, SOCK_STREAM, 0);
    struct sockaddr_in server;
    server.sin_addr.s_addr = inet_addr(server_ip.c_str());
    server.sin_family = AF_INET;
    server.sin_port = htons(port);

    /* Connect to the remote server */
    if (connect(client_fd, (struct sockaddr *)&server, sizeof(server)) < 0)
    {
        perror("Connection failed");
        return 1;
    }

    /* Set the socket to non-blocking mode */
    int flags = fcntl(client_fd, F_GETFL, 0);
    fcntl(client_fd, F_SETFL, flags | O_NONBLOCK);

    /* Return the socket */
    return client_fd;
};


int create_tcp_server_socket(int port) {
    /* Local variables */
    struct sockaddr_in server_addr;

    /* Create the TCP server-side socket */
    int server_fd = socket(AF_INET, SOCK_STREAM, 0);
    if (server_fd < 0)
    {
        exit_program_with_code(EXIT_FAILURE, "failed at creating tcp server socket");
    }

    /* Set socket options */
    int flags = fcntl(server_fd, F_GETFL, 0);
    if (flags == -1) {
        release_fd(server_fd);
        exit_program_with_code(EXIT_FAILURE, "failed at getting socket flags -1");
    }
    flags |= O_NONBLOCK;
    if (fcntl(server_fd, F_SETFL, flags) == -1) {
        release_fd(server_fd);
        exit_program_with_code(EXIT_FAILURE, "failed at getting socket flags as NONBLOCK");
    }

    /* Bind the socket */
    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_addr.s_addr = INADDR_ANY;
    server_addr.sin_port = htons(port);
    if (bind(server_fd, (struct sockaddr *) &server_addr, sizeof(server_addr)) < 0) {
        release_fd(server_fd);
        exit_program_with_code(1, "failed at binding socket");
    }

    /* Listen for connections */
    if (listen(server_fd, 5) < 0) {
        release_fd(server_fd);
        exit_program_with_code(1, "failed at listening on socket");
    }

    /* Return the server socket */
    return server_fd;
}

void handle_signal(int sig)
{
    fprintf(stderr, "Caught signal %d\n", sig);
    exit(0);
}


int create_netlink_socket(int group_id) {
    int socket_fd;
    struct sockaddr_nl process_addr; // Addresses for recv, send, bind
    int result;

    /* Set up the socket */
    socket_fd = socket(PF_NETLINK, SOCK_DGRAM, NETLINK_CONNECTOR);
    if (socket_fd == -1)
        exit_program_with_code(EXIT_FAILURE, "socket");

    /* Initialize the address structs */
    memset(&process_addr, 0, sizeof(struct sockaddr_nl));
    process_addr.nl_family = AF_NETLINK;
    process_addr.nl_pid = getpid(); // This process' PID
    if (group_id == -1)
        process_addr.nl_groups = CN_IDX_IWLAGN;
    else
        process_addr.nl_groups = group_id;

    /* Bind the socket */
    if (bind(socket_fd, (struct sockaddr *)&process_addr, sizeof(struct sockaddr_nl)) == -1) {
        release_fd(socket_fd);
        exit_program_with_code(EXIT_FAILURE, "bind");
    }

    /* Subscribe to the netlink group */
    __u32 on = process_addr.nl_groups;
    result = setsockopt(socket_fd, 270, NETLINK_ADD_MEMBERSHIP, &on, sizeof(on));
    if (result) {
        release_fd(socket_fd);
        exit_program_with_code(EXIT_FAILURE, "setsockopt");
    }

    /* Set up the "handle_signal" function as this program's signal handler */
    signal(SIGINT, handle_signal);

    return socket_fd;
}


int exit_program_with_code(int code, std::string msg) {
    if (code == EXIT_SUCCESS) {
        exit(EXIT_SUCCESS);
    } else {
        perror(msg.c_str());
        exit(code);
    }
};


int release_fd(int& fd) {
    if (fd != -1) {
        close(fd);
        fd = -1;
    }
}